library(rmarkdown)    # You need this library to run this template.
library(epuRate) 
#library(tidyverse)
library(tidyr)
library(dplyr)
library(cowplot)
library(ggplot2)
library(readr)
library(wesanderson)
#library(rjson)
library(RJSONIO)
source("R_rainclouds.R") # for the raincloud plot

line.plot <- function(data, dv, iv, grp, ylims, cols){
          ggplot(data, aes_string(x=iv, y=dv, col=grp)) +
            geom_line(aes_string(group=grp), size=1.1) + geom_point() +
            facet_wrap(~sub, nrow=1) +  
            scale_fill_manual(values=cols) +
            scale_color_manual(values=cols) + 
            ylab(dv) + xlab(iv) + ylim(ylims) +
            theme(panel.border = element_blank(), 
                  panel.grid.major =   element_blank(),
            panel.grid.minor = element_blank(), 
            axis.line = element_line(colour = "black"))
}

get_participant_data <- function(subjects, sessions, data_path) {
  # this function loads each participant's data
  # and concatenates them into a longform dataset
  fn <- 'sub-0%d_ses-%d_task-learn-att-v1-test-v1_events.tsv'
  get_subject_strings <- function(i) {
    get_session_strings <- function(j) dir(sprintf(paste(data_path, "sub-0%d_ses-%d_task-learn-att-v1", sep = "/"), i, j), pattern=sprintf(fn, i, j), full.names = TRUE)
    do.call(cbind, lapply(sessions, get_session_strings))
  }
  files <- do.call(rbind, lapply(subjects, get_subject_strings))
  rownames(files) <- subjects
  colnames(files) <- sessions
  resplog <- function(i, j) read.table(files[as.character(i),as.character(j)], sep = "\t", header = TRUE)
  d <- do.call(rbind, lapply(subjects, function(i) do.call(rbind, lapply(sessions, function (j) resplog(i, j)))))
  
  # get trials to getthe reward condition 
  fn <- 'sub-0%d_ses-%d_task-learn-att-v1-test-v1_trls.csv'
  files <- do.call(rbind, lapply(subjects, get_subject_strings))      
  rownames(files) <- subjects
  colnames(files) <- sessions
  eventlog <- function(i, j) {
    e = read.table(files[as.character(i),as.character(j)], sep = ",", header = TRUE)
    e$sub = i
    e$sess= j
    e
  }
  e <- do.call(rbind, lapply(subjects, function(i) do.call(rbind, lapply(sessions, function (j) eventlog(i, j)))))
  names(e)[1] = "t"
  e <- e[, -c(2, 4:8 )]  
  d <- inner_join(d, e, by=c("sub", "sess", "t"))
  # ALLOCATE CUES
  # 1 = left arrow, 2 = right arrow, 3 = bidirectional arrow
  d$cert <- NA
  d$cert[ d$loc == 1 & d$cue == 1 ] = ".8"
  d$cert[ d$loc == 2 & d$cue == 1 ] = ".2"
  d$cert[ d$loc == 2 & d$cue == 2 ] = ".8"
  d$cert[ d$loc == 1 & d$cue == 2 ] = ".2"  
  d$cert[ d$cue == 3 ] = ".5"
  d$cert <- as.factor(d$cert)
  
  d$loc <- as.factor(d$loc)
  levels(d$loc) <- c("left", "right")
  d$sub <- as.factor(d$sub)
  d$rew <- as.factor(d$rew)
  levels(d$rew) <- c("0", "50")
  d$reward_type <- as.factor(d$reward_type)
  levels(d$reward_type) <- c("htgt/hdst", "htgt/ldst", "ltgt/ldst", "ltgt/hdst")
  
  d <- d %>% filter(t > 120) # get rid of the first block of trials
  
  d
}

subjects = c(1, 2, 3)
sessions = 1
data_path = "~/Dropbox/MC-Projects/imaging-value-cert-att/task-code/pilot_data_analysis/pilot_expCues"
raw.data <- get_participant_data(subjects, sessions, data_path)

# Show it:
raw.data %>% head(5)

get.sub.data <- function(subject, data, sd_reject=2.5, RT_min=0.1){
  # this function takes the subject index (defined in the variable subject)
  # and all the raw data
  # it filters the data to get the individual subject's data, then trims to
  # get the correct RTs that are > .1 s, and are < 3 * sd from the median
  # for each certainty and reward condition
  sub.data <- data %>% filter(sub == subject)
  sub.data <- sub.data %>% 
                filter(rt > RT_min) %>%
                filter(resp == 1) %>%
                group_by(cert, reward_type) %>%
                filter(rt < median(rt) + sd_reject*sd(rt)) 
}
# apply the function defined above across subjects, put
# results into a single dataframe
subjects = c("1", "2", "3")  
sub.data <- lapply(subjects, get.sub.data, data = raw.data)
sub.data <- do.call(rbind, sub.data)